package com.hyz.yuaiagent.chatmemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * 基于文件持久化的对话记忆实现
 * 使用Kryo序列化技术将对话记录持久化到本地文件系统
 * 线程不安全实现，适用于单线程场景
 */
public class FileBasedChatMemory implements ChatMemory {

    private static final Kryo kryo = new Kryo();

    static {
        // 配置Kryo序列化器
        // 禁用强制注册机制以提高兼容性
        // 使用标准实例化策略支持更多类型反序列化
        kryo.setRegistrationRequired(false);
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    private final String BASE_DIR;

    /**
     * 构造文件对话记忆实例
     *
     * @param dir 文件存储基础目录路径
     */
    public FileBasedChatMemory(String dir) {
        // 设置临时目录路径为当前用户目录下的tmp子目录
        this.BASE_DIR = dir;
        File baseDir = new File(dir);
        if (!baseDir.exists()) {
            baseDir.mkdirs();
        }
    }


    /**
     * 添加消息到指定对话
     *
     * @param conversationId 对话唯一标识
     * @param messages       待添加的消息列表
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        List<Message> conversationMessages = getOrCreateConversation(conversationId);
        conversationMessages.addAll(messages);
        saveConversation(conversationId, conversationMessages);
    }

    /**
     * 获取对话的最近N条消息
     *
     * @param conversationId 对话唯一标识
     * @param lastN          需要获取的最近消息数量
     * @return 包含最近N条消息的列表（按原始顺序）
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Message> allMessages = getOrCreateConversation(conversationId);
        return allMessages.stream()
                .skip(Math.max(0, allMessages.size() - lastN))
                .toList();
    }

    /**
     * 清除指定对话的所有记录
     *
     * @param conversationId 需要清除的对话标识
     */
    @Override
    public void clear(String conversationId) {
        File file = getConversationFile(conversationId);
        if (file.exists()) {
            file.delete();
        }
    }

    /**
     * 获取或创建对话消息列表
     *
     * @param conversationId 对话标识
     * @return 反序列化后的消息列表
     */
    private List<Message> getOrCreateConversation(String conversationId) {
        File file = getConversationFile(conversationId);
        List<Message> messages = new ArrayList<>();
        if (file.exists()) {
            try (Input input = new Input(new FileInputStream(file))) {
                messages = kryo.readObject(input, ArrayList.class);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return messages;
    }

    /**
     * 持久化对话消息到文件
     *
     * @param conversationId 对话标识
     * @param messages       需要持久化的消息列表
     */
    private void saveConversation(String conversationId, List<Message> messages) {
        File file = getConversationFile(conversationId);
        try (Output output = new Output(new FileOutputStream(file))) {
            kryo.writeObject(output, messages);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 获取对话对应的存储文件
     *
     * @param conversationId 对话标识
     * @return 对应的文件对象
     */
    private File getConversationFile(String conversationId) {
        return new File(BASE_DIR, conversationId + ".kryo");
    }
}
